import streamlit as st
import cv2
import torch
import os
import time
from ultralytics import YOLO
from utils.model_utils import find_latest_best_model, clean_old_files

def main():
    st.title("单SAR图识别")

    # 上传图像
    uploaded_file = st.file_uploader("选择SAR图像...", type=["png", "jpg", "jpeg"])
    if uploaded_file is not None:
        try:
            # 保存上传的图像到本地
            timestamp = int(time.time())
            image_filename = f"{timestamp}_{uploaded_file.name}"
            image_path = os.path.join("uploads", image_filename)
            os.makedirs(os.path.dirname(image_path), exist_ok=True)
            with open(image_path, "wb") as f:
                f.write(uploaded_file.read())

            # 读取并显示图像
            image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
            if image is None:
                st.error("无法读取图像，请确保上传的文件是有效的图像格式。")
            else:
                st.image(image, caption="上传的SAR图像", width=400)

                # 选择模型
                # model_options = ["YOLOv11n", "YOLOv11s", "YOLOv11m", "YOLOv11l", "YOLOv11x"]
                # 选择模型
                model_options = {
                    "YOLOv11n": "yolo11n_best.pt",
                    "YOLOv11s": "yolo11s_best.pt",
                    "YOLOv11m": "yolo11m_best.pt",
                    "YOLOv11l": "yolo11l_best.pt",
                    "YOLOv11x": "yolo11x_best.pt"
                }
                selected_model_name = st.selectbox("选择YOLOv11模型", list(model_options.keys()))
				# 设置置信度阈值
                confidence_threshold = st.slider("置信度阈值", 0.0, 1.0, 0.5, 0.05)

                if st.button("开始检测"):
                    # 查找最新的 best 模型文件
                    current_dir = os.getcwd()
                    parent_dir = os.path.dirname(current_dir)
                    models_dir = os.path.join(parent_dir, "models", "v1")
                    selected_model_path = os.path.join(models_dir, model_options[selected_model_name])
                    if not os.path.exists(selected_model_path):
                        selected_model_path = find_latest_best_model(models_dir)
                        st.warning(f"找不到选定的模型文件，使用最新的最佳模型: {selected_model_path}")

                    st.write(f"选择的模型文件: {selected_model_path}")
                    # 创建保存检测结果的目录
                    save_dir = os.path.join(current_dir, "runs", "detect", "predict")
                    if not os.path.exists(save_dir):
                        os.makedirs(save_dir)

                    # 加载模型
                    model = YOLO(selected_model_path)

                    # 使用模型预测
                    with st.spinner("正在检测..."):
                        results = model.predict(source=image_path, save=True, imgsz=900, conf=confidence_threshold, save_dir=save_dir)

                    # 显示检测结果
                    st.subheader("检测结果")
                    for result in results:
                        for box in result.boxes:
                            if box.conf >= confidence_threshold:
                                st.write(f"检测到 {result.names[int(box.cls)]} 位置: {box.xyxy[0].tolist()}, 置信度: {box.conf.item():.2f}")

                    # 获取预测结果的保存路径
                    if results and hasattr(results[0], 'save_dir'):
                        # 提取原始文件名
                        original_image_filename = os.path.basename(image_path)
                        # 构建保存的文件名
                        saved_image_filename = os.path.splitext(original_image_filename)[0] + '.jpg'  # 假设保存为.jpg格式
                        # 构建完整的保存路径
                        predicted_image_path = os.path.join(results[0].save_dir, saved_image_filename)

                        # 如果文件不存在，可能是由于YOLO添加了不同的后缀，尝试查找可能的文件
                        if not os.path.exists(predicted_image_path):
                            # 尝试查找带有常见后缀的文件
                            for ext in ['.jpg', '.png']:
                                temp_path = os.path.splitext(predicted_image_path)[0] + ext
                                if os.path.exists(temp_path):
                                    predicted_image_path = temp_path
                                    break

                        # 读取预测结果图像
                        if os.path.exists(predicted_image_path):
                            predicted_image = cv2.imread(predicted_image_path)
                            if predicted_image is not None:
                                # 将BGR图像转换为RGB
                                predicted_image_rgb = cv2.cvtColor(predicted_image, cv2.COLOR_BGR2RGB)
                                st.image(predicted_image_rgb, caption="YOLO预测的结果图像", width=400)
                            else:
                                st.error("未能加载预测结果图像。")
                        else:
                            st.error("没有找到预测结果图像。")
                    else:
                        st.error("没有找到预测结果。")

                    # 清理20次以前的临时文件
                    clean_old_files("uploads", max_keep=20)

        except Exception as e:
            st.error(f"发生错误: {e}")

if __name__ == "__main__":
    main()